# -*- coding: utf-8 -*-
import numpy as np
from gensim.models import Word2Vec
from tqdm import tqdm

# 获取所有词表 word2vec.model.wv.index2word
class CustomWord2Vec:
    """
    基于gensim的可增量以及可分批训练的word2vec
    """
    def __init__(self, size=32, window=5, min_count=1, workers=20, sg=0, load_path=None):
        self.size = size
        self.window = window
        self.min_count = min_count
        self.workers = workers
        self.sg = sg
        self.load_path = load_path
        self.build_model()

    def build_model(self):
        if self.load_path:
            self.model = Word2Vec.load(self.load_path)
        else:
            self.model = Word2Vec(min_count=self.min_count,
                                  size=self.size,
                                  window=self.window,
                                  workers=self.workers,
                                  sg=self.sg,
                                  compute_loss=True)

    def build_vocab(self, sentences, update):
        self.model.build_vocab(sentences, update=update)

    def train(self, sentences, epochs):
        self.model.train(sentences, total_examples=self.model.corpus_count, epochs=epochs)

    def fit(self, sentences, update_vocab=True, epochs=1):
        if update_vocab:
            try:
                self.build_vocab(sentences, update=update_vocab)
            except Exception as e:
                print(e)
                self.build_vocab(sentences, update=False)
        self.train(sentences, epochs)
        print("loss:", self.model.get_latest_training_loss())

    def fit_generator(self, generator, update_vocab=True, epochs=1):
        for sentences in tqdm(generator):
            self.fit(sentences, update_vocab, epochs)

    def save(self, model_path):
        self.model.save(model_path)

    def get(self, key, default=None):
        if key in self.model.wv:
            return self.model.wv[key]
        if default is None:
            return np.zeros(self.size)
        return default


